#include "sessionmanager.h"
// Qt核心头文件
#include <QObject>
#include <QTimer>
#include <QDebug>
#include <QRegularExpression>
#include <QDateTime>
#include <QString>
#include <QList>
#include <QMap>
#include <QVector>
#include <QPair>

// 项目头文件
#include "common.h"
#include "session.h"
#include "sessionmanager.h"

SessionManager::SessionManager(QObject *parent) : QObject(parent)
{
    m_sessionTimeoutMs = 30000; // 默认30秒超时
    m_timeoutTimerId = -1;
}

SessionManager::~SessionManager()
{
    stopSessionTimeoutCheck();
    m_sessions.clear();
}

// 处理捕获的数据包，尝试将会话相关的数据包分组
void SessionManager::processPacket(const PacketInfo &packet)
{
    // 识别应用层协议类型
    SessionType type = identifyApplicationProtocol(packet);
    
    // 对于不支持的协议类型，不创建会话
    if (type == SessionType::OTHER) {
        return;
    }
    
    // 创建或更新会话
    QString sessionId = createOrUpdateSession(packet, type);
    
    // 发送会话更新信号
    if (m_sessions.contains(sessionId)) {
        emit sessionUpdated(m_sessions[sessionId]);
        
        // 如果是新会话（数据包索引列表只有一个元素），发送新会话信号
        if (m_sessions[sessionId].packetIndices.size() == 1) {
            emit newSessionDetected(m_sessions[sessionId]);
        }
    }
}

// 获取所有会话
QList<ProtocolSession> SessionManager::getAllSessions() const
{
    return m_sessions.values();
}

// 根据ID获取会话
ProtocolSession *SessionManager::getSessionById(const QString &sessionId)
{
    if (m_sessions.contains(sessionId)) {
        return &m_sessions[sessionId];
    }
    return nullptr;
}

// 根据数据包索引获取相关会话
QList<ProtocolSession *> SessionManager::getSessionsByPacketIndex(int packetIndex)
{
    QList<ProtocolSession *> result;
    
    for (auto &session : m_sessions) {
        if (session.packetIndices.contains(packetIndex)) {
            result.append(&session);
        }
    }
    
    return result;
}

// 获取指定类型的会话
QList<ProtocolSession *> SessionManager::getSessionsByType(SessionType type)
{
    QList<ProtocolSession *> result;
    
    for (auto &session : m_sessions) {
        if (session.type == type) {
            result.append(&session);
        }
    }
    
    return result;
}

// 清除所有会话
void SessionManager::clearAllSessions()
{
    m_sessions.clear();
}

// 获取会话总数
int SessionManager::getSessionCount() const
{
    return m_sessions.size();
}

// 开始会话超时检查
void SessionManager::startSessionTimeoutCheck(int intervalMs)
{
    // 停止之前的定时器
    if (m_timeoutTimerId != -1) {
        killTimer(m_timeoutTimerId);
        m_timeoutTimerId = -1;
    }
    
    // 启动新的定时器
    m_timeoutTimerId = startTimer(intervalMs);
}

// 停止会话超时检查
void SessionManager::stopSessionTimeoutCheck()
{
    if (m_timeoutTimerId != -1) {
        killTimer(m_timeoutTimerId);
        m_timeoutTimerId = -1;
    }
}

// 设置会话超时时间
void SessionManager::setSessionTimeout(int timeoutMs)
{
    m_sessionTimeoutMs = timeoutMs;
}

// 检查会话超时
void SessionManager::checkSessionTimeouts()
{
    QDateTime now = QDateTime::currentDateTime();
    QList<QString> toRemove;
    
    for (auto it = m_sessions.begin(); it != m_sessions.end(); ++it) {
        ProtocolSession &session = it.value();
        
        // 只检查活跃状态的会话
        if (session.state == SessionState::ACTIVE) {
            // 计算会话最后活动时间（使用最后一个数据包的时间戳）
            qint64 timeSinceLastActivity = session.startTime.msecsTo(now);
            
            if (timeSinceLastActivity > m_sessionTimeoutMs) {
                // 将会话标记为超时
                session.updateState(SessionState::TIMED_OUT);
                emit sessionUpdated(session);
            }
        }
    }
}

// 事件过滤器，处理定时器事件
void SessionManager::timerEvent(QTimerEvent *event)
{
    if (event->timerId() == m_timeoutTimerId) {
        checkSessionTimeouts();
    }
    
    QObject::timerEvent(event);
}

// 从数据包信息中提取应用层协议类型
SessionType SessionManager::identifyApplicationProtocol(const PacketInfo &packet)
{
    QString appLayerDetail = packet.applicationLayerDetail.toLower();
    
    if (appLayerDetail.contains("http") && !appLayerDetail.contains("https")) {
        return SessionType::HTTP;
    } else if (appLayerDetail.contains("https")) {
        return SessionType::HTTPS;
    } else if (appLayerDetail.contains("dns")) {
        return SessionType::DNS;
    } else if (appLayerDetail.contains("ftp")) {
        return SessionType::FTP;
    } else if (appLayerDetail.contains("ssh")) {
        return SessionType::SSH;
    } else if (appLayerDetail.contains("smtp")) {
        return SessionType::SMTP;
    } else if (appLayerDetail.contains("pop3")) {
        return SessionType::POP3;
    } else if (appLayerDetail.contains("imap")) {
        return SessionType::IMAP;
    }
    
    return SessionType::OTHER;
}

// 为数据包创建或更新会话
QString SessionManager::createOrUpdateSession(const PacketInfo &packet, SessionType type)
{
    int srcPort = 0, dstPort = 0;
    extractPorts(packet, srcPort, dstPort);
    
    // 尝试找到匹配的现有会话
    for (auto &session : m_sessions) {
        if (isSameSession(packet, session)) {
            // 添加数据包到现有会话
            session.addPacket(packet.index);
            session.totalBytes += packet.length;
            
            // 更新会话的结束时间
            session.endTime = QDateTime::currentDateTime();
            
            // 根据协议类型解析特定信息
            switch (type) {
            case SessionType::HTTP:
                parseHTTP(packet, session);
                break;
            case SessionType::DNS:
                parseDNS(packet, session);
                break;
            case SessionType::FTP:
                parseFTP(packet, session);
                break;
            default:
                break;
            }
            
            return session.id;
        }
    }
    
    // 如果没有找到匹配的会话，创建新会话
    ProtocolSession newSession;
    newSession.type = type;
    newSession.typeString = sessionTypeToString(type);
    newSession.clientAddress = packet.sourceIp;
    newSession.serverAddress = packet.destIp;
    newSession.clientPort = srcPort;
    newSession.serverPort = dstPort;
    newSession.addPacket(packet.index);
    newSession.totalBytes = packet.length;
    
    // 根据协议类型解析特定信息
    switch (type) {
    case SessionType::HTTP:
        parseHTTP(packet, newSession);
        break;
    case SessionType::DNS:
        parseDNS(packet, newSession);
        break;
    case SessionType::FTP:
        parseFTP(packet, newSession);
        break;
    default:
        break;
    }
    
    // 将新会话添加到会话列表
    m_sessions[newSession.id] = newSession;
    
    return newSession.id;
}

// 检查两个数据包是否属于同一个会话
bool SessionManager::isSameSession(const PacketInfo &packet, const ProtocolSession &session)
{
    // 对于TCP/UDP连接，使用四元组(IP+端口组合)来标识会话
    // 允许双向匹配（请求和响应可能交换IP和端口）
    bool directMatch = (packet.sourceIp == session.clientAddress && 
                       packet.destIp == session.serverAddress);
                        
    bool reverseMatch = (packet.sourceIp == session.serverAddress && 
                        packet.destIp == session.clientAddress);
                        
    // 尝试从数据包中提取端口信息
    int srcPort = 0, dstPort = 0;
    if (extractPorts(packet, srcPort, dstPort)) {
        directMatch = directMatch && (srcPort == session.clientPort) && (dstPort == session.serverPort);
        reverseMatch = reverseMatch && (srcPort == session.serverPort) && (dstPort == session.clientPort);
    }
                        
    // 对于相同的协议类型，且IP和端口匹配（直接或反向），认为是同一会话
    SessionType packetType = identifyApplicationProtocol(packet);
    
    return (packetType == session.type) && (directMatch || reverseMatch);
}

// 解析HTTP协议，提取请求和响应信息
void SessionManager::parseHTTP(const PacketInfo &packet, ProtocolSession &session)
{
    QString appLayerDetail = packet.applicationLayerDetail;
    
    // 检查是否包含HTTP请求方法
    if (appLayerDetail.contains("GET ") || 
        appLayerDetail.contains("POST ") || 
        appLayerDetail.contains("PUT ") || 
        appLayerDetail.contains("DELETE ") || 
        appLayerDetail.contains("HEAD ") || 
        appLayerDetail.contains("OPTIONS ")) {
        
        // 这是一个HTTP请求
        session.requestCount++;
        
        // 尝试提取URL或路径
        QRegularExpression regex("(GET|POST|PUT|DELETE|HEAD|OPTIONS) +([^ ]+)");
        QRegularExpressionMatch match = regex.match(appLayerDetail);
        if (match.hasMatch()) {
            QString method = match.captured(1);
            QString path = match.captured(2);
            
            // 如果是第一个请求，设置会话摘要
            if (session.summary.isEmpty()) {
                session.summary = QString("%1 %2").arg(method).arg(path);
            }
            
            // 存储请求信息
            session.additionalInfo[QString("Request%1").arg(session.requestCount)] = 
                QString("%1 %2").arg(method).arg(path);
        }
        
    } else if (appLayerDetail.contains("HTTP/") && 
               (appLayerDetail.contains(" 200 ") || 
                appLayerDetail.contains(" 404 ") || 
                appLayerDetail.contains(" 500 "))) {
        
        // 这是一个HTTP响应
        session.responseCount++;
        
        // 尝试提取状态码
        QRegularExpression regex("HTTP/[\\d.]+\\s+(\\d+)\\s+([^\\n]+)");
        QRegularExpressionMatch match = regex.match(appLayerDetail);
        if (match.hasMatch()) {
            QString statusCode = match.captured(1);
            QString statusText = match.captured(2);
            
            // 存储响应信息
            session.additionalInfo[QString("Response%1").arg(session.responseCount)] = 
                QString("%1 %2").arg(statusCode).arg(statusText);
            
            // 如果是200 OK响应，可以考虑将会话标记为完成
            if (statusCode == "200") {
                session.updateState(SessionState::COMPLETED);
                emit sessionCompleted(session);
            }
        }
    }
}

// 解析DNS协议，提取查询和响应信息
void SessionManager::parseDNS(const PacketInfo &packet, ProtocolSession &session)
{
    QString appLayerDetail = packet.applicationLayerDetail;
    
    // 这里简化实现，实际应该解析DNS头部和记录
    // 对于这个示例，我们假设应用层详情中包含了足够的信息
    
    if (appLayerDetail.contains("Query")) {
        // 这是一个DNS查询
        session.requestCount++;
        
        // 尝试提取查询域名
        QRegularExpression regex("Query: *([^ ]+)");
        QRegularExpressionMatch match = regex.match(appLayerDetail);
        if (match.hasMatch()) {
            QString domain = match.captured(1);
            
            if (session.summary.isEmpty()) {
                session.summary = QString("Query: %1").arg(domain);
            }
            
            session.additionalInfo[QString("Query%1").arg(session.requestCount)] = domain;
        }
        
    } else if (appLayerDetail.contains("Response")) {
        // 这是一个DNS响应
        session.responseCount++;
        
        // 尝试提取响应记录
        QRegularExpression regex("Response: *([^ ]+)");
        QRegularExpressionMatch match = regex.match(appLayerDetail);
        if (match.hasMatch()) {
            QString response = match.captured(1);
            
            session.additionalInfo[QString("Response%1").arg(session.responseCount)] = response;
            
            // DNS响应通常表示会话完成
            session.updateState(SessionState::COMPLETED);
            emit sessionCompleted(session);
        }
    }
}

// 解析FTP协议，提取命令和响应信息
void SessionManager::parseFTP(const PacketInfo &packet, ProtocolSession &session)
{
    QString appLayerDetail = packet.applicationLayerDetail;
    
    // FTP命令通常是大写的，以换行符结尾
    QRegularExpression cmdRegex("^([A-Z]+)\\s*([^\\n]*)");
    QRegularExpressionMatch cmdMatch = cmdRegex.match(appLayerDetail);
    
    if (cmdMatch.hasMatch()) {
        // 这是一个FTP命令
        session.requestCount++;
        
        QString cmd = cmdMatch.captured(1);
        QString params = cmdMatch.captured(2);
        
        if (session.summary.isEmpty()) {
            session.summary = QString("FTP Command: %1").arg(cmd);
        }
        
        session.additionalInfo[QString("Command%1").arg(session.requestCount)] = 
            QString("%1 %2").arg(cmd).arg(params);
        
    } else if (appLayerDetail.startsWith("220 ") || 
               appLayerDetail.startsWith("250 ") || 
               appLayerDetail.startsWith("331 ") || 
               appLayerDetail.startsWith("530 ")) {
        // 这是一个FTP响应
        session.responseCount++;
        
        // 提取状态码和响应文本
        QRegularExpression respRegex("^(\\d{3})\\s+([^\\n]*)");
        QRegularExpressionMatch respMatch = respRegex.match(appLayerDetail);
        
        if (respMatch.hasMatch()) {
            QString statusCode = respMatch.captured(1);
            QString respText = respMatch.captured(2);
            
            session.additionalInfo[QString("Response%1").arg(session.responseCount)] = 
                QString("%1 %2").arg(statusCode).arg(respText);
        }
    }
}

// 从传输层详情中提取端口信息
bool SessionManager::extractPorts(const PacketInfo &packet, int &srcPort, int &dstPort)
{
    QString transportDetail = packet.transportLayerDetail;
    
    // 尝试从传输层详情中提取源端口和目标端口
    QRegularExpression srcPortRegex("源端口: *([0-9]+)");
    QRegularExpression dstPortRegex("目的端口: *([0-9]+)");
    
    QRegularExpressionMatch srcMatch = srcPortRegex.match(transportDetail);
    QRegularExpressionMatch dstMatch = dstPortRegex.match(transportDetail);
    
    if (srcMatch.hasMatch() && dstMatch.hasMatch()) {
        srcPort = srcMatch.captured(1).toInt();
        dstPort = dstMatch.captured(1).toInt();
        return true;
    }
    
    return false;
}